from rest_framework import viewsets, status
from rest_framework.decorators import action
from rest_framework.response import Response
from rest_framework.permissions import IsAuthenticated
from .models import Stock, UserStock
from .serializers import StockSerializer, StockScreenerSerializer
import pandas as pd
from .utils import TradingStrategy
from .utils.stock_data import StockDataService

class StockViewSet(viewsets.ModelViewSet):
    queryset = Stock.objects.all()
    serializer_class = StockSerializer
    permission_classes = [IsAuthenticated]

    @action(detail=True, methods=['post'])
    def follow(self, request, pk=None):
        stock = self.get_object()
        user = request.user
        
        if UserStock.objects.filter(user=user, stock=stock).exists():
            return Response({"detail": "已经关注过该股票"}, status=status.HTTP_400_BAD_REQUEST)
        
        UserStock.objects.create(user=user, stock=stock)
        return Response({"detail": "关注成功"}, status=status.HTTP_201_CREATED)

    @action(detail=True, methods=['post'])
    def unfollow(self, request, pk=None):
        stock = self.get_object()
        user = request.user
        
        try:
            user_stock = UserStock.objects.get(user=user, stock=stock)
            user_stock.delete()
            return Response({"detail": "取消关注成功"}, status=status.HTTP_200_OK)
        except UserStock.DoesNotExist:
            return Response({"detail": "未关注该股票"}, status=status.HTTP_400_BAD_REQUEST)


    @action(detail=True, methods=['get'])
    def trading_suggestion(self, request, pk=None):
        """获取股票交易建议"""
        stock = self.get_object()
        strategy = TradingStrategy(stock)
        suggestion = strategy.get_trading_suggestion()
        return Response(suggestion)

    @action(detail=False, methods=['get'])
    def all_suggestions(self, request):
        """获取所有关注股票的交易建议"""
        user = request.user
        user_stocks = UserStock.objects.filter(user=user).select_related('stock')
        suggestions = []
        
        for user_stock in user_stocks:
            strategy = TradingStrategy(user_stock.stock)
            suggestion = strategy.get_trading_suggestion()
            suggestions.append(suggestion)
            
        return Response(suggestions)

    @action(detail=False, methods=['post'])
    def screen(self, request):
        """
        股票筛选接口
        最多返回市值最大的50只股票
        """
        serializer = StockScreenerSerializer(data=request.data)
        if not serializer.is_valid():
            return Response(serializer.errors, status=status.HTTP_400_BAD_REQUEST)
        
        data_service = StockDataService()
        df = data_service.get_stock_basic_info(
            industry=serializer.validated_data.get('industry'),
            market_value_min=serializer.validated_data.get('market_value_min'),
            market_value_max=serializer.validated_data.get('market_value_max'),
            turnover_rate_min=serializer.validated_data.get('turnover_rate_min'),
            pe_min=serializer.validated_data.get('pe_min'),
            pe_max=serializer.validated_data.get('pe_max')
        )
        
        if df is None:
            return Response({"detail": "获取数据失败"}, status=status.HTTP_500_INTERNAL_SERVER_ERROR)
        
        # 按市值降序排序并限制返回数量
        df = df.sort_values('total_mv', ascending=False).head(50)
        
        # 转换为列表格式
        results = []
        for _, row in df.iterrows():
            results.append({
                'code': row['ts_code'],
                'name': row['name'],
                'industry': row['industry'],
                'market_value': float(row['total_mv']) / 100000,  # 转换为亿元
                'turnover_rate': float(row['turnover_rate']),
                'pe_ratio': float(row['pe']) if not pd.isna(row['pe']) else None,
                'list_date': row['list_date']
            })
        
        return Response({
            'count': len(results),
            'total_count': len(df),  # 添加符合条件的总数
            'results': results,
            'message': '最多返回市值最大的50只股票'
        }) 